import json
from tasks.bvd.graph_class import *

# NOTE: An error in jsons:  edges['src_addr'] is src_id, not src_addr

# class to transfer json to graph
class json_constructor:
    def __init__(self, json_name, graph_type):
        self.name = json_name
        self.type = graph_type

    def construct_by_inst(self, max_node_num=1000, max_edge_num=10000):
        # construct node/edge at asm instrtuction level
        # each node is specified by its address
        self.max_node_num = max_node_num
        self.max_edge_num = max_edge_num
        with open(self.name, 'r') as f:
            data = json.load(f)
        if self.type == 'cfg':
            ret = inst_graph('cfg')
            ret.file_name = self.name
            bb_nodes = data['nodes']
            bb_edges = data['edges']
            # step1: construct nodes && inner basic block edges
            dict_id_first_inst_idx = {}
            for bb_struct in bb_nodes:
                is_first_inst = True
                first_inst_idx = -1
                
                previous_inst_idx=-1
                
                for inst in bb_struct['disasm']:
                    if ret.get_node_num() >= self.max_node_num:
                        break
                    if is_first_inst:
                        node = inst_node(inst['address'], inst['mnemonic'], inst['op_str'], is_first_inst, bb_struct['id'])
                        is_first_inst = False
                        ret.inst_nodes.append(node)
                        first_inst_idx = ret.get_node_num()-1
                        previous_inst_idx=first_inst_idx
                        dict_id_first_inst_idx[bb_struct['id']] = first_inst_idx
                    else:
                        node = inst_node(inst['address'], inst['mnemonic'], inst['op_str'], is_first_inst, bb_struct['id'])
                        ret.inst_nodes.append(node)
                        if ret.get_edge_num() < self.max_edge_num:
                            # temporarily undirected graph
                            edge_1 = inst_edge(first_inst_idx, ret.get_node_num()-1)
                            edge_2 = inst_edge(ret.get_node_num()-1, first_inst_idx)
                            
                            # neighor edge
                            edge_3 = inst_edge(previous_inst_idx, ret.get_node_num()-1)
                            edge_4 = inst_edge(ret.get_node_num()-1, previous_inst_idx)
                            
                            ret.inst_edges.append(edge_1)
                            ret.inst_edges.append(edge_2)
                            ret.inst_edges.append(edge_3)
                            ret.inst_edges.append(edge_4)
                            
                            previous_inst_idx=ret.get_node_num()-1
                            
            # step2: construct cross basic block edges
            for bb_edge in bb_edges:
                if ret.get_edge_num() < self.max_edge_num:
                    if bb_edge['src_id'] in dict_id_first_inst_idx.keys() and bb_edge['dst_id'] in dict_id_first_inst_idx.keys():
                        edge1 = inst_edge(dict_id_first_inst_idx[bb_edge['src_id']], dict_id_first_inst_idx[bb_edge['dst_id']])
                        ret.inst_edges.append(edge1)
                        edge2 = inst_edge(dict_id_first_inst_idx[bb_edge['dst_id']], dict_id_first_inst_idx[bb_edge['src_id']])
                        ret.inst_edges.append(edge2)
            return ret
        else:
            pass #TODO: support other types of graph
    
# class to transfer former APIs to graph
class api_constructor:
    def __init__(self, graph_type, graph_data):
        self.type = graph_type
        self.data = graph_data

    def construct_by_inst(self, max_node_num=1000, max_edge_num=10000):
        # construct node/edge at asm instrtuction level
        # each node is specified by its address
        self.max_node_num = max_node_num
        self.max_edge_num = max_edge_num
        if self.type == 'cfg':
            ret = inst_graph('cfg')
            bb_nodes = self.data['nodes']
            bb_edges = self.data['edges']
            # step1: construct nodes && inner basic block edges
            dict_id_first_inst_idx = {}
            for bb_struct in bb_nodes:
                is_first_inst = True
                first_inst_idx = -1
                for inst in bb_struct['disasm']:
                    if ret.get_node_num() >= self.max_node_num:
                        break
                    if is_first_inst:
                        node = inst_node(inst['address'], inst['mnemonic'], inst['op_str'], is_first_inst, bb_struct['id'])
                        is_first_inst = False
                        ret.inst_nodes.append(node)
                        first_inst_idx = ret.get_node_num()-1
                        dict_id_first_inst_idx[bb_struct['id']] = first_inst_idx
                    else:
                        node = inst_node(inst['address'], inst['mnemonic'], inst['op_str'], is_first_inst, bb_struct['id'])
                        ret.inst_nodes.append(node)
                        if ret.get_edge_num() < self.max_edge_num:
                            # temporarily undirected graph
                            edge_1 = inst_edge(first_inst_idx, ret.get_node_num()-1)
                            edge_2 = inst_edge(ret.get_node_num()-1, first_inst_idx)
                            ret.inst_edges.append(edge_1)
                            ret.inst_edges.append(edge_2)
            # step2: construct cross basic block edges
            for bb_edge in bb_edges:
                if ret.get_edge_num() < self.max_edge_num:
                    edge1 = inst_edge(dict_id_first_inst_idx[bb_edge['src_addr']], dict_id_first_inst_idx[bb_edge['dst_addr']])
                    ret.inst_edges.append(edge1)
                    edge2 = inst_edge(dict_id_first_inst_idx[bb_edge['dst_addr']], dict_id_first_inst_idx[bb_edge['src_addr']])
                    ret.inst_edges.append(edge2)
            return ret
        else:
            pass #TODO: support other types of graph

# class to transfer graph to dataload
class dataset_constructor:
    def __init__(self, graph):
        self.graph = graph
    
    def split_dataset():
        pass
    def construct_dataset():
        pass
    
def test():
    pass


if __name__ == '__main__':
    test()